#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI 批量下载插件
基于 downlist.py 功能
"""

import os
import sys
import time
import urllib.parse
import requests
import re
from pathlib import Path
from typing import List, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed

# 导入统一的环境配置
from public.env_setup import setup_environment, ENV_CONFIG, BATCHDOWN_DIR
from public.download_utils import download_file_with_hf

# 使用统一的环境配置
ENV_CONFIG = setup_environment()

class BatchDownloader:
    """批量下载器插件类"""
    
    # 插件排序
    order = 3
    
    def __init__(self):
        """初始化下载器"""
        # 延迟获取环境配置
        self.env_config = ENV_CONFIG
        self.workspace_path = self.env_config['PROJECT_ROOT']
        # 获取 py_scripts 目录
        self.config_file = self.env_config['CONFIG_DIR'] / "model_urls.conf"
        self.save_path = BATCHDOWN_DIR
        self.default_filename = ""
        
        # 创建下载目录
        self.save_path.mkdir(exist_ok=True)
        
        # 颜色定义
        self.RED = "\033[31m"
        self.GREEN = "\033[32m"
        self.YELLOW = "\033[33m"
        self.BLUE = "\033[34m"
        self.CYAN = "\033[36m"
        self.PURPLE = "\033[35m"
        self.RESET = "\033[0m"
    
    def start(self) -> bool:
        """启动批量下载器"""
        print(f"\n{self.PURPLE}■■■ 批量下载器 v1.0（ComfyUI插件版） ■■■{self.RESET}")
        print(f"配置文件: {self.CYAN}{self.config_file}{self.RESET}")
        print(f"下载目录: {self.CYAN}{self.save_path.absolute()}{self.RESET}")
        
        print("0. 退出")
        print("1. 下载全部模型")
        print("2. 下载指定模型")
        print("3. 显示模型列表")
        choice = input("请选择功能 (0-3): ").strip()
        if choice == "0" or choice.lower() == "q":
            print("已退出插件。")
            return False

        urls = self.read_urls_from_file()
        if not urls:
            print(f"{self.YELLOW}⚠️ 配置文件中没有有效 URL{self.RESET}")
            return False

        print(f"\n{self.BLUE}找到 {len(urls)} 个下载任务{self.RESET}")
        
        # 确认下载
        confirm = input(f"\n是否开始批量下载？(y/N): ").strip().lower()
        if confirm != 'y':
            print(f"{self.YELLOW}操作已取消{self.RESET}")
            return False

        # 使用4个线程并发下载
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = [
                executor.submit(self.download_task, url, idx+1, len(urls))
                for idx, url in enumerate(urls)
            ]
            
            success_count = 0
            for future in as_completed(futures):
                url, success = future.result()
                if success:
                    success_count += 1
        
        print(f"\n{self.GREEN}✅ 批量下载完成！成功: {success_count}/{len(urls)}{self.RESET}")
        return True
    
    @staticmethod
    def Launch(self):
        """静态启动方法"""
        downloader = BatchDownloader()
        return downloader.start()
    
    def read_urls_from_file(self) -> List[str]:
        """从配置文件读取URL列表"""
        urls = []
        if not self.config_file.exists():
            print(f"{self.RED}❌ 找不到配置文件: {self.config_file}{self.RESET}")
            return urls

        with open(self.config_file, "r", encoding="utf-8-sig") as f:  # utf-8-sig 可处理 BOM
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):  # 跳过空行和注释
                    continue
                # 可选：移除行内注释（# 之后的内容）
                if "#" in line:
                    line = line.split("#", 1)[0].strip()
                if line:
                    urls.append(line)
        return urls

    def transform_huggingface_url(self, url: str) -> str:
        """转换HuggingFace URL为镜像地址"""
        if url.startswith("https://huggingface.co/") or url.startswith("https://hf-mirror.com/"):
            parsed = urllib.parse.urlparse(url)
            parts = parsed.path.strip("/").split("/")
            if len(parts) > 2:
                parts[2] = "resolve"

            # 替换 netloc 并重组路径
            new_path = "/" + "/".join(parts)
            new_parsed = parsed._replace(netloc="hf-mirror.com", path=new_path)
            return urllib.parse.urlunparse(new_parsed)
        return url

    def extract_filename_from_url(self, url: str) -> str:
        """从URL中提取文件名"""
        parsed_url = urllib.parse.urlparse(url)
        path = parsed_url.path
        basename = os.path.basename(path)
        filename = urllib.parse.unquote(basename)

        if '.' in filename and not filename.startswith('.'):
            return filename
        else:
            return f"{filename}.download"

    def download_file_with_hf(self, url: str, filename: str = "", save_path: str = "") -> Tuple[bool, str]:
        """下载文件（支持HuggingFace）"""
        url = self.transform_huggingface_url(url)
        auto_filename = self.extract_filename_from_url(url)

        if not filename:
            filename = auto_filename
        elif '.' not in filename:
            ext = os.path.splitext(auto_filename)[1]
            filename += ext

        filepath = self.save_path / filename

        print(f"{self.CYAN}下载: {url}{self.RESET}")
        print(f"{self.CYAN}保存: {filepath}{self.RESET}")

        try:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
            }
            
            with requests.get(url, stream=True, timeout=30, headers=headers) as r:
                r.raise_for_status()
                total_size = int(r.headers.get('content-length', 0))
                
                with open(filepath, 'wb') as f:
                    downloaded = 0
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            downloaded += len(chunk)
                            if total_size > 0:
                                percent = (downloaded / total_size) * 100
                                print(f"\r{self.CYAN}进度: {percent:.1f}%{self.RESET}", end='', flush=True)
                    
                    print()  # 换行
            
            return True, str(filepath)
            
        except Exception as e:
            print(f"{self.RED}下载失败: {e}{self.RESET}")
            if filepath.exists():
                filepath.unlink()
            return False, str(e)

    def download_task(self, url: str, idx: int, total: int) -> Tuple[str, bool]:
        """单个下载任务"""
        print(f"\n{self.BLUE}=== [{idx}/{total}] 开始下载: {url}{self.RESET}")
        success, result = self.download_file_with_hf(url=url, filename=self.default_filename, save_path=str(self.save_path))
        if success:
            print(f"{self.GREEN}✅ 文件已保存: {result}{self.RESET}")
            return url, True
        else:
            print(f"{self.RED}❌ 下载失败: {url}, 错误: {result}{self.RESET}")
            return url, False

def main():
    """主函数"""
    downloader = BatchDownloader()
    downloader.start()

if __name__ == "__main__":
    main() 